#!/usr/bin/env python3
# 导入必要的库
import argparse
import os
import sys
import time
import grpc
import traceback
import threading

# 添加P4Runtime工具库路径
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 './utils/')
)
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.switch import ShutdownAllSwitchConnections

import struct


H2_MAC="08:00:00:00:02:02"
H2_IP="10.0.2.2"
H1_MAC="08:00:00:00:01:01"
H1_IP="10.0.1.1"

S1_P2="00:00:00:00:01:02"
S1_P3="00:00:00:00:01:03"
S1_P4="00:00:00:00:01:04"
S1_P5="00:00:00:00:01:05"

S6_P2="00:00:00:00:06:02"
S6_P3="00:00:00:00:06:03"
S6_P4="00:00:00:00:06:04"
S6_P5="00:00:00:00:06:05"


def writeIpv4_lpmRules(p4info_helper, sw, dst_ip, dst_mac, port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, 32)
        },
        action_name="MyIngress.ipv4_forward",
        action_params={
            "dstAddr": dst_mac,
            "port"   : port,
        })
    sw.WriteTableEntry(table_entry)

def writeEcmp_groupRules(p4info_helper, sw, dst_ip, base, count):
    table_entry =p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_group",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip, 32)
        },
        action_name="MyIngress.set_ecmp_select",
        action_params={
            "ecmp_base":base,
            "ecmp_count":count,
        })
    sw.WriteTableEntry(table_entry)

def writeEcmp_nhopRules(p4info_helper, sw, select, nhop_dmac, nhop_ipv4, port):
    table_entry=p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_nhop",
        match_fields={
            "meta.ecmp_select":select
        },
        action_name="MyIngress.set_nhop",
        action_params={
            "nhop_dmac":nhop_dmac,
            "nhop_ipv4":nhop_ipv4,
            "port":port,
        })
    sw.WriteTableEntry(table_entry)

def writeMyEgressRules(p4info_helper, sw, port, smac):
    table_entry=p4info_helper.buildTableEntry(
        table_name="MyEgress.send_frame",
        match_fields={
            "standard_metadata.egress_port":port
        },
        action_name="MyEgress.rewrite_mac",
        action_params={
            "smac":smac,
        })
    sw.WriteTableEntry(table_entry)

def printGrpcError(e):
    """打印gRPC通信错误的详细信息"""
    print("gRPC Error:", e.details(), end=' ')
    status_code = e.code()
    print(f"({status_code.name})", end=' ')
    traceback_obj = sys.exc_info()[2]
    print(f"[{traceback_obj.tb_frame.f_code.co_filename}:{traceback_obj.tb_lineno}]")






def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 创建与交换机s1的连接
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s1',
            address='127.0.0.1:50051',
            device_id=0,
            proto_dump_file='logs/s1-p4runtime-requests.txt')

        # 创建与交换机s2的连接
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s2',
            address='127.0.0.1:50052',
            device_id=1,
            proto_dump_file='logs/s2-p4runtime-requests.txt')

        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s3',
            address='127.0.0.1:50053',
            device_id=2,
            proto_dump_file='logs/s3-p4runtime-requests.txt')

        s4 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s4',
            address='127.0.0.1:50054',
            device_id=3,
            proto_dump_file='logs/s4-p4runtime-requests.txt')

        s5 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s5',
            address='127.0.0.1:50055',
            device_id=4,
            proto_dump_file='logs/s5-p4runtime-requests.txt')

        s6 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name='s6',
            address='127.0.0.1:50056',
            device_id=5,
            proto_dump_file='logs/s6-p4runtime-requests.txt')



        # 发送主控制器仲裁消息
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()
        s4.MasterArbitrationUpdate()
        s5.MasterArbitrationUpdate()
        s6.MasterArbitrationUpdate()

        # 在交换机上安装P4程序
        s1.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s1")

        s2.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s2")

        s3.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s3")

        s4.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s4")

        s5.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s5")

        s6.SetForwardingPipelineConfig(p4info=p4info_helper.p4info,
                                       bmv2_json_file_path=bmv2_file_path)
        print("Installed P4 Program on s6")

        writeEcmp_groupRules(p4info_helper,s1,"10.0.0.1",0,4)
        writeEcmp_nhopRules(p4info_helper, s1, 0, S1_P2, H2_IP, 2)
        writeEcmp_nhopRules(p4info_helper, s1, 1, S1_P3, H2_IP, 3)
        writeEcmp_nhopRules(p4info_helper, s1, 2, S1_P4, H2_IP, 4)
        writeEcmp_nhopRules(p4info_helper, s1, 3, S1_P5, H2_IP, 5)
        writeIpv4_lpmRules(p4info_helper, s1, H1_IP, H1_MAC, 1)

        writeEcmp_groupRules(p4info_helper,s6,"10.0.0.2",0,4)
        writeEcmp_nhopRules(p4info_helper, s6, 0, S6_P2, H1_IP, 2)
        writeEcmp_nhopRules(p4info_helper, s6, 1, S6_P3, H1_IP, 3)
        writeEcmp_nhopRules(p4info_helper, s6, 2, S6_P4, H1_IP, 4)
        writeEcmp_nhopRules(p4info_helper, s6, 3, S6_P5, H1_IP, 5)
        writeIpv4_lpmRules(p4info_helper, s6, H2_IP, H2_MAC, 1)

        writeIpv4_lpmRules(p4info_helper, s2, H1_IP, H1_MAC, 1)
        writeIpv4_lpmRules(p4info_helper, s2, H2_IP, H2_MAC, 2)
        writeIpv4_lpmRules(p4info_helper, s3, H1_IP, H1_MAC, 1)
        writeIpv4_lpmRules(p4info_helper, s3, H2_IP, H2_MAC, 2)
        writeIpv4_lpmRules(p4info_helper, s4, H1_IP, H1_MAC, 1)
        writeIpv4_lpmRules(p4info_helper, s4, H2_IP, H2_MAC, 2)
        writeIpv4_lpmRules(p4info_helper, s5, H1_IP, H1_MAC, 1)
        writeIpv4_lpmRules(p4info_helper, s5, H2_IP, H2_MAC, 2)

        writeMyEgressRules(p4info_helper, s1, 2, S1_P2)
        writeMyEgressRules(p4info_helper, s1, 3, S1_P3)
        writeMyEgressRules(p4info_helper, s1, 4, S1_P4)
        writeMyEgressRules(p4info_helper, s1, 5, S1_P5)

        writeMyEgressRules(p4info_helper, s6, 2, S6_P2)
        writeMyEgressRules(p4info_helper, s6, 3, S6_P3)
        writeMyEgressRules(p4info_helper, s6, 4, S6_P4)
        writeMyEgressRules(p4info_helper, s6, 5, S6_P5)







    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)
    except Exception as e:
        print(f"Unexpected error: {str(e)}")
        traceback.print_exc()

    ShutdownAllSwitchConnections()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print(f"\np4info file not found: {args.p4info}\nHave you run 'make'?")
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print(f"\nBMv2 JSON file not found: {args.bmv2_json}\nHave you run 'make'?")
        parser.exit(1)

    main(args.p4info, args.bmv2_json)
